package com.ks.core.api;

import com.google.common.util.concurrent.RateLimiter;
import com.ks.core.api.anno.RateLimit;
import com.ks.core.api.config.ApiConfig;
import com.ks.core.core.Result;
import com.ks.core.core.ResultCode;
import com.ks.core.exception.GunsException;
import com.ks.core.util.JsonUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.lang.reflect.Method;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * @author kingdee
 */
@Slf4j
public class RateLimiterInterceptor implements HandlerInterceptor {

    private Map<Method, RateLimiter> rateLimiterMap = new ConcurrentHashMap<>(16);

    private ApiConfig apiConfig;

    public RateLimiterInterceptor(ApiConfig apiConfig) {
        this.apiConfig = apiConfig;
    }

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        log.info("进入 [RateLimiterInterceptor] 拦截器");
        if (!(handler instanceof HandlerMethod)) {
            log.info("不是 [HandlerMethod] 类型，放行");
            return true;
        }
        HandlerMethod handlerMethod = (HandlerMethod) handler;
        Method method = handlerMethod.getMethod();
        if (!method.isAnnotationPresent(RateLimit.class)) {
            log.info("方法 [" + method.getName() + "] 不包含 [@RateLimite] 放行");
            return true;
        }
        RateLimiter rateLimiter;
        if ((rateLimiter = rateLimiterMap.get(method)) == null) {
            double rate = method.getAnnotation(RateLimit.class).value();
            if (rate <= 0) {
                rate = apiConfig.getDefaultCallRate();
            }
            rateLimiter = RateLimiter.create(rate);
            rateLimiterMap.put(method, rateLimiter);
        }
        method.getAnnotation(RateLimit.class).value();
        long timeOutMill = apiConfig.getTimeOutMill();
        if (rateLimiter.tryAcquire(timeOutMill, TimeUnit.MILLISECONDS)) {
            log.info("获取到令牌，放行");
            return true;
        }
        log.info(timeOutMill + "ms未获取到令牌，拦截");
        response.setContentType(MediaType.APPLICATION_JSON_UTF8_VALUE);
        response.getWriter().write(JsonUtil.objectToJsonStr(
                Result.fail(new GunsException(ResultCode.FAIL.getCode(), "调用频率过快，请稍后再试"))));
        return false;
    }
}
